using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.Globalization;
using System.IO;

using Atomic.Libraries;
using Atomic.Libraries.Mathematics;
using Atomic.Phonons;
using Atomic.Structures;
using Atomic.Vasp;
using Atomic.Vasp.LinearResponse;
using Atomic.Thermodynamics.Phonons;
using Atomic.Thermodynamics;
using Atomic.Libraries.Plotting.Functions;
using System.Linq;

namespace Atomic.Interfaces.Phonopy
{
	public static class PhonopyPhonon
	{
		public static PhononStateDensity CreateStateDensity(IVaspLinearResponseResult result)
		{
			// Maybe a higher default value?
			// http://phonopy.sourceforge.net/examples.html
			return CreateStateDensity(result, 31);
		}

		public static PhononStateDensity CreateStateDensity(IVaspLinearResponseResult result, int mesh)
		{
			VaspStructure structure = result.FinalStructure;

			using (TempDirectory temp = new TempDirectory())
			{
				WriteStateDensityConfiguration(result.FinalStructure, mesh, temp.Directory);
				WriteForceConstants(result, temp.Directory);
				VaspPoscar.Write(structure.PrimitiveStructure, temp.Directory.AddFile("POSCAR-unitcell"));

				Process p = new Process();
				p.StartInfo.FileName = "phonopy";
				p.StartInfo.Arguments = "--dos -c POSCAR-unitcell pdos.conf";
				p.StartInfo.WorkingDirectory = temp.Directory.FullName;
				p.StartInfo.UseShellExecute = false;
				p.StartInfo.RedirectStandardOutput = true;
				p.StartInfo.RedirectStandardError = true;
				p.Start();
				p.WaitForExit();

				if (p.ExitCode != 0)
				{
					string s = p.StandardError.ReadToEnd().Trim();
					if (s == "")
					{
						// Sometimes Phonopy writes errors to standard output.
						s = p.StandardOutput.ReadToEnd().Trim();
					}

					throw new Exception(string.Format("Failed running phonopy. Error message: {0}", s));
				}

				return ParseStateDensity(temp.Directory, 3 * result.FinalStructure.PrimitiveStructure.Sites.Count);
			}
		}

		/*public static PhononDispersion CreateDispersion(VaspLinearResponseResult result, IEnumerable<Tuple<SpaceVector, string>> points)
		{
		}*/

		public static PhononDispersion CreateDispersion(IVaspLinearResponseResult result, IEnumerable<SpaceVector> points)
		{
			return CreateDispersion(result, new List<SpaceVector>(points).ToArray());
		}

		private static PhononDispersion CreateDispersion(IVaspLinearResponseResult result, SpaceVector[] points)
		{
			VaspStructure structure = result.FinalStructure;

			using (TempDirectory temp = new TempDirectory())
			{
				WriteDispersionConfiguration(structure, points, temp.Directory);
				WriteForceConstants(result, temp.Directory);
				VaspPoscar.Write(structure.PrimitiveStructure, temp.Directory.AddFile("POSCAR-unitcell"));

				Process p = new Process();
				p.StartInfo.FileName = "phonopy";
				p.StartInfo.Arguments = "-c POSCAR-unitcell band.conf";
				p.StartInfo.WorkingDirectory = temp.Directory.FullName;
				p.StartInfo.UseShellExecute = false;
				p.StartInfo.RedirectStandardOutput = true;
				p.StartInfo.RedirectStandardError = true;
				p.Start();
				p.WaitForExit();

				if (p.ExitCode != 0)
				{
					throw new Exception(string.Format("Failed running phonopy. Error message: {0}", p.StandardError.ReadToEnd().Trim()));
				}

				return ParseDispersion(result, temp.Directory, points);
			}
		}

		private static void WriteStateDensityConfiguration(VaspStructure structure, int mesh, DirectoryInfo directory)
		{
			using (StreamWriter writer = new StreamWriter(directory.AddFile("pdos.conf").FullName))
			{
				WriteBaseConfiguration(structure, writer);

				writer.WriteLine("MP = {0} {0} {0}", mesh);
				//writer.WriteLine("GAMMA_CENTER = .TRUE.");
			}
		}

		private static void WriteDispersionConfiguration(VaspStructure structure, SpaceVector[] points, DirectoryInfo directory)
		{
			using (StreamWriter writer = new StreamWriter(directory.AddFile("band.conf").FullName))
			{
				WriteBaseConfiguration(structure, writer);

				// http://phonopy.sourceforge.net/setting-tags.html#band-band-points
				writer.Write("BAND = ");
				for (int i = 0; i < points.Length; i++)
				{
					if (i > 0)
					{
						writer.Write("  ");
					}
					writer.Write(string.Format(CultureInfo.InvariantCulture, "{0} {1} {2}", points[i].X, points[i].Y, points[i].Z));
				}
				writer.WriteLine();

				writer.WriteLine("BAND_POINTS = 101");
			}
		}

		private static void WriteBaseConfiguration(VaspStructure structure, StreamWriter writer)
		{
			// The supercell relative to the lattice vectors of the primitive cell. 
			// http://phonopy.sourceforge.net/setting-tags.html#dim
			SpaceMatrix dim = structure.PrimitiveStructure.Lattice.InversePhysicalMap * structure.Lattice.PhysicalMap;
			writer.Write("DIM =");
			for (int i = 0; i < 3; i++)
			{
				if (i > 0)
				{
					writer.Write(" ");
				}
				for (int j = 0; j < 3; j++)
				{
					double d = dim[i, j];
					if (Math.Abs(d - Math.Round(d)) > 1.0e-4)
					{
						throw new Exception("Invalid scaling of primitive lattice; numbers must be integers.");
					}

					writer.Write(" {0}", (int)Math.Round(d));
				}
			}
			writer.WriteLine();

			/*// Atoms are used to create the dynamical matrix, which is weighted by atomic masses.
			writer.Write("ATOM_NAME =");
			foreach (Site site in structure.PrimitiveStructure.Sites)
			{
				writer.Write(" {0}", site.Atom.Symbol);
			}
			writer.WriteLine();*/

			// Atomic masses are used to create the dynamical matrix, which is weighted by atomic masses.
			// http://phonopy.sourceforge.net/setting-tags.html#mass
			writer.Write("MASS =");
			foreach (Site site in structure.PrimitiveStructure.Sites)
			{
				writer.Write(" {0}", site.Atom.Mass);
			}
			writer.WriteLine();

			// http://phonopy.sourceforge.net/setting-tags.html#id6
			writer.WriteLine("FORCE_CONSTANTS = READ");
		}

		private static void WriteForceConstants(IVaspLinearResponseResult result, DirectoryInfo directory)
		{
			// File format here:
			// http://phonopy.sourceforge.net/input-files.html#format

			using (StreamWriter writer = new StreamWriter(directory.AddFile("FORCE_CONSTANTS").FullName))
			{
				int n = result.ForceConstants.Count;

				writer.WriteLine("{0}", n);
				for (int i = 0; i < n; i++)
				{
					for (int j = 0; j < n; j++)
					{
						writer.WriteLine("{0} {1}", i + 1, j + 1);
						for (int k = 0; k < 3; k++)
						{
							for (int l = 0; l < 3; l++)
							{
								writer.Write(string.Format(CultureInfo.InvariantCulture, " {0}", result.ForceConstants[i, j][k, l]));
							}
							writer.WriteLine();
						}
					}
				}
			}
		}

		private static PhononStateDensity ParseStateDensity(DirectoryInfo directory, int phononStates)
		{
			List<PhononStateDensityPoint> points = new List<PhononStateDensityPoint>();
			using (StreamReader sr = new StreamReader(directory.AddFile("total_dos.dat").FullName))
			{
				// Ignore first header line.
				sr.ReadLine();

				string s;
				while ((s = sr.ReadLine()) != null)
				{
					s = s.Trim();
					while (s.Contains("  "))
					{
						s = s.Replace("  ", " ");
					}
					double f = double.Parse(s.Split(' ')[0], CultureInfo.InvariantCulture);
					double d = double.Parse(s.Split(' ')[1], CultureInfo.InvariantCulture);

					// Ignore leading zero density points; but include if points are already added.
					if (d != 0.0 || points.Count > 0)
					{
						points.Add(new PhononStateDensityPoint(f, d));
					}
				}
			}

			return new PhononStateDensity(points, phononStates);
		}

		private static PhononDispersion ParseDispersion(IVaspLinearResponseResult result, DirectoryInfo directory, SpaceVector[] symmetryPoints)
		{
			// The number of phonon branches is three times the number of atoms in the primitive cell.
			int m = 3 * result.FinalStructure.PrimitiveStructure.Sites.Count;

			List<List<PhononDispersionPoint>> points = new List<List<PhononDispersionPoint>>();
			for (int i = 0; i < m; i++)
			{
				points.Add(new List<PhononDispersionPoint>());
			}

			using (StreamReader reader = new StreamReader(directory.AddFile("band.yaml").FullName))
			{
				SpaceVector q = null;
				double d = double.NaN;
				List<double> f = new List<double>();

				string s;
				while ((s = reader.ReadLine()) != null)
				{
					if (s.Contains("q-position:"))
					{
						q = AtomicEnvironment.ParseSpaceVector(s.Substring(s.IndexOf("[") + 1, s.IndexOf("]") - s.IndexOf("[") - 1));
					}
					if (s.Contains("distance:"))
					{
						d = double.Parse(s.Substring(s.IndexOf(":") + 1).Trim(), CultureInfo.InvariantCulture);
					}
					if (s.Contains("frequency:"))
					{
						if (q == null)
						{
							throw new Exception("Unexcepted number of frequencies.");
						}

						f.Add(double.Parse(s.Substring(s.IndexOf(":") + 1).Trim(), CultureInfo.InvariantCulture));
					}
					if (f.Count == m)
					{
						for (int i = 0; i < m; i++)
						{
							points[i].Add(new PhononDispersionPoint(new Position(result.FinalStructure.PrimitiveStructure.Lattice, q), d, f[i]));
						}

						q = null;
						d = double.NaN;
						f.Clear();
					}
				}
			}

			List<PhononDispersionBranch> branches = new List<PhononDispersionBranch>();
			for (int i = 0; i < m; i++)
			{
				branches.Add(new PhononDispersionBranch(points[i]));
			}

			return new PhononDispersion(branches, symmetryPoints);
		}
		
		public static IPotential CreatePhonopyPotential(IVaspLinearResponseResult result, int mesh)
		{
			//Structure s0 = result.FinalStructure.PrimitiveStructure;
			VaspStructure structure = result.FinalStructure;
			//LatticeRepeatMatrix m = LatticeRepeatMatrix.Create(s0, structure);

			// Check by a double size input cell. Should double the per unit cell values.

			// In the Phonopy output 1 mol is defined as N_A unit cells.
			// 1 kJ/"mol" = 0.01036410 eV/cell, 1 J/mol = 0.001 kJ/mol
			double conv = 0.01036410;
			double conv2 = 0.001 * conv;

			using (TempDirectory temp = new TempDirectory())
			{
				using (StreamWriter writer = new StreamWriter(temp.Directory.AddFile("mesh.conf").FullName))
				{
					WriteBaseConfiguration(structure, writer);

					writer.WriteLine("MP = {0} {0} {0}", mesh);
					//writer.WriteLine("GAMMA_CENTER = .TRUE.");
				}
				WriteForceConstants(result, temp.Directory);
				VaspPoscar.Write(structure.PrimitiveStructure, temp.Directory.AddFile("POSCAR"));

				//File.WriteAllLines(temp.Directory.AddFile("mesh.conf").FullName,
				//	new string[] { "ATOM_NAME = " + FormatAtoms(s0.Sites), "DIM = " + PhonopyProblem.FormatRepeatMatrix(m), string.Format("MP = {0} {0} {0}", mesh), "TMIN = 0", "TMAX = 1002", "TSTEP = 3" });
				//File.WriteAllBytes(temp.Directory.AddFile("FORCE_SETS").FullName, result.Force.GetData());
				//File.WriteAllBytes(temp.Directory.AddFile("POSCAR").FullName, Problem.Poscar.GetData());

				Process p = new Process();
				p.StartInfo.FileName = "phonopy";
				p.StartInfo.Arguments = "-t mesh.conf";
				p.StartInfo.WorkingDirectory = temp.Directory.FullName;
				p.StartInfo.UseShellExecute = false;
				p.StartInfo.RedirectStandardOutput = true;
				p.StartInfo.RedirectStandardError = true;

				p.Start();
				
				List<Tuple<double, double, double, double>> thermalProperties = new List<Tuple<double, double, double, double>>();
				bool reading = false;
				string line;
				while ((line = p.StandardOutput.ReadLine()) != null)
				{
					if (line.StartsWith("#"))
					{
						reading = true;
					}
					else if (reading)
					{
						line = line.Trim();
						while (line.Contains("  "))
						{
							line = line.Replace("  ", " ");
						}

						string[] cols = line.Split(' ');
						if (cols.Length == 5)
						{
							double t = double.Parse(cols[0].Trim(), CultureInfo.InvariantCulture);
							double f = conv * double.Parse(cols[1].Trim(), CultureInfo.InvariantCulture);
							double s = conv2 * double.Parse(cols[2].Trim(), CultureInfo.InvariantCulture);
							double cv = conv2 * double.Parse(cols[3].Trim(), CultureInfo.InvariantCulture);

							thermalProperties.Add(Tuple.Create<double, double, double, double>(t, f, s, cv));
						}
						else
						{
							reading = false;
						}
					}
				}

				p.WaitForExit();

				if (p.ExitCode != 0)
				{
					throw new Exception(string.Format("Running phonopy failed: {0}", p.StandardError.ReadToEnd()));
				}

				return new PhonopyPotential(
					new LinearInterpolationFunction(thermalProperties.Select(tp => new PlaneVector(tp.Item1, tp.Item2))),
					new LinearInterpolationFunction(thermalProperties.Select(tp => new PlaneVector(tp.Item1, tp.Item3))),
					new LinearInterpolationFunction(thermalProperties.Select(tp => new PlaneVector(tp.Item1, tp.Item4))));
			}
		}

		private static string FormatAtoms(IEnumerable<Site> sites)
		{
			List<Atom> atoms = new List<Atom>();
			List<string> symbols = new List<string>();
			foreach (Site site in sites)
			{
				if (!atoms.Contains(site.Atom))
				{
					atoms.Add(site.Atom);
					symbols.Add(site.Atom.Symbol);
				}
			}

			return string.Join(" ", symbols);
		}

		[Serializable]
		private class PhonopyPotential : IPotential
		{
			private IPlotFunction freeEnergy, entropy, heatCapacity;

			public PhonopyPotential(IPlotFunction freeEnergy, IPlotFunction entropy, IPlotFunction heatCapacity)
			{
				this.freeEnergy = freeEnergy;
				this.entropy = entropy;
				this.heatCapacity = heatCapacity;
			}

			public double FreeEnergy(double temperature)
			{
				return freeEnergy.Value(temperature);
			}

			public double Entropy(double temperature)
			{
				return entropy.Value(temperature);
			}

			public double HeatCapacity(double temperature)
			{
				return heatCapacity.Value(temperature);
			}
		}
	}
}
